import torch
from sklearn.tree import DecisionTreeRegressor
from sklearn.metrics import mean_squared_error
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split

# 生成数据
X, y = make_regression(n_samples=1000, n_features=5, noise=0.1, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 参数
n_trees = 50  # 多少棵树
learning_rate = 0.1  # 学习率

# 转换数据为 PyTorch 张量
X_train_torch = torch.tensor(X_train, dtype=torch.float32)
y_train_torch = torch.tensor(y_train, dtype=torch.float32)

# 初始化预测值
y_pred_train = torch.zeros_like(y_train_torch)

# 训练 GBDT
trees = []
for i in range(n_trees):
    # 计算梯度（残差）
    residuals = y_train_torch - y_pred_train

    # 用决策树拟合梯度
    tree = DecisionTreeRegressor(max_depth=3)
    tree.fit(X_train, residuals.numpy())
    trees.append(tree)

    # 更新预测值
    y_pred_train += learning_rate * torch.tensor(tree.predict(X_train), dtype=torch.float32)

    # 计算损失
    mse = mean_squared_error(y_train, y_pred_train.numpy())
    print(f"Iteration {i+1}: MSE = {mse:.4f}")
